/*
 * @brief:
 *
 * @author: dariaye@tencent.com
 * @date: 2012-6-12
 * @version: 1.0
 * @changes:
 *    <author>        <time>        <version >        <desc>
 *     dariaye             2012-6-12              1.0                 Create
 */

#include <stdlib.h>
#include <stdarg.h>
#include <unistd.h>
#include <assert.h>
#include <mpi.h>
#include <sched.h>

#include "util.h"

#define kBufferSize 8192

WALLE_BEGIN

void FPrintf(MPI_Comm comm, FILE* fp, const char format[],...)
{
    int rank;

    if (!comm) comm = MPI_COMM_WORLD;
    MPI_Comm_rank(comm,&rank);
    if(rank == 0)
    {
        va_list argp;
        char buffer[kBufferSize];
        va_start(argp, format);
        vsnprintf(buffer, sizeof(buffer), format, argp);
        fprintf(fp, "%s", buffer);
        va_end(argp);
    }
    MPI_Barrier(comm);
}

void SyncPrintf(MPI_Comm comm,const char format[],...)
{
    int info, size, rank;
    info = MPI_Comm_size(comm, &size);
    CHKERRQ(info);
    info = MPI_Comm_rank(comm, &rank);
    CHKERRQ(info);

    va_list argp;
    char buffer[kBufferSize];
    va_start(argp, format);
    vsnprintf(buffer, sizeof(buffer), format, argp);
    va_end(argp);

    for(int i = 0; i < size; ++i)
    {
        if(i == rank)
        {
            printf("######## Rank %d ########\n", rank);
            printf("%s", buffer);
        }

        MPI_Barrier(comm);
    }
}

void Printf(MPI_Comm comm,const char format[],...)
{
    int rank;

    if (!comm) comm = MPI_COMM_WORLD;
    MPI_Comm_rank(comm,&rank);
    if(rank == 0)
    {
        va_list argp;
        char buffer[kBufferSize];
        va_start(argp, format);
        vsnprintf(buffer, sizeof(buffer), format, argp);
        printf("%s", buffer);
        va_end(argp);
    }
    MPI_Barrier(comm);
}

void SetAffinity(int cpu_id)
{
#ifdef LIB_01
    cpu_set_t cpu_mask;
    CPU_ZERO(&cpu_mask);
    CPU_SET(cpu_id, &cpu_mask);

    sched_setaffinity(0, &cpu_mask);
#endif

#ifdef LIB_02
    unsigned long cpu_mask = (unsigned long) 1;
    cpu_mask=cpu_mask<<cpu_id;
    sched_setaffinity(0, sizeof(unsigned long), &cpu_mask);
#endif

#ifdef LIB_03
    cpu_set_t cpu_mask;
    CPU_ZERO(&cpu_mask);
    CPU_SET(cpu_id, &cpu_mask);

    sched_setaffinity(0, sizeof(cpu_set_t), &cpu_mask);
#endif
}

WALLE_END
